package com.example.llollama.controller;

import com.example.llollama.commponent.MyHtmlReader;
import com.example.llollama.commponent.MyTikaDocumentReader;

import java.net.MalformedURLException;
import java.util.List;
import java.util.Map;

import com.example.llollama.util.PromptUtil;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.http.MediaType;
import org.springframework.util.CollectionUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

/**
 * @Description: TODO
 * @author: lonely
 * @date: 2024年10月18日 23:40
 */
@RequestMapping("/rag")
@RestController
@RequiredArgsConstructor
@Slf4j
public class RAGController {
    private final VectorStore vectorStore;
    private final OllamaChatModel chatModel;
    private final MyTikaDocumentReader docReader;
    private final MyHtmlReader htmlReader;

    @GetMapping("/create")
    public String ragCreate() {
        //1.提取文本内容
        String filePath = "刘磊简历.txt";
        TextReader textReader = new TextReader(filePath);
        textReader.getCustomMetadata().put("filePath", filePath);
        List<Document> documents = textReader.get();
        log.info("文档分割之前:{}", documents.size());

        //2.文本切分段落
        storeDocuments(documents);
        return "创建完成！！！";
    }

    @GetMapping("/createPdf")
    public String createPdf() {

        List<Document> documents = getDocsFromPdf();
        log.info("文档分割之前:{}", documents.size());

        //2.文本切分段落
        storeDocuments(documents);
        return "创建完成！！！";
    }

    private void storeDocuments(List<Document> documents) {
        TokenTextSplitter splitter = new TokenTextSplitter(1200, 350, 5, 100, true);
        System.out.println("开始调用");
        List<Document> splitDocuments = splitter.apply(documents);
        System.out.println("调用完成");
        log.info("文档分割之后：{}", splitDocuments.size());

        //3.段落写入向量数据库
        vectorStore.add(splitDocuments);
    }

    @GetMapping("/createDoc")
    public String createDoc() {
        List<Document> documents = docReader.loadText();
        //2.文本切分段落
        storeDocuments(documents);
        return "创建完成！！！";
    }

    @GetMapping("/createHtml")
    public String createHtml(@RequestParam String url) throws MalformedURLException {
        //String url="https://baike.baidu.com/item/%E9%BB%84%E5%AE%B6%E9%A9%B9/77734";
        List<Document> documents = htmlReader.loadText(url);
        PromptUtil.setSource(documents,url);
        storeDocuments(documents);
        return "创建完成！！！";
    }

    List<Document> getDocsFromPdf() {

        PagePdfDocumentReader pdfReader = new PagePdfDocumentReader("classpath:/yjl.pdf",
            PdfDocumentReaderConfig.builder()
                .withPageTopMargin(0)
                .withPageExtractedTextFormatter(ExtractedTextFormatter.builder()
                    .withNumberOfTopTextLinesToDelete(0)
                    .build())
                .withPagesPerDocument(1)
                .build());
        List<Document> read = pdfReader.read();
        return read.stream().map(e -> {
            Map<String, Object> metadata = e.getMetadata();
            metadata.put("filePath", "yjl.pdf");
            return new Document(e.getContent(), metadata);
        }).toList();
    }

    @GetMapping(value="/search")
    public Flux<String> search(@RequestParam("query") String query) {
        long currentTime1 = System.currentTimeMillis();
        //检索挂载信息
        List<Document> documents = vectorStore.similaritySearch(query);
        String fromFiles = "";

        if(!CollectionUtils.isEmpty(documents)){
            fromFiles=documents.get(0).getMetadata().get("filePath").toString();
        }
        //提取相关信息
        String info = "";
        if (documents.size() > 0) {
            info = String.join(",", documents.stream().map(Document::getContent).toList());
        }
        log.info("检索耗时：{}ms", System.currentTimeMillis() - currentTime1);
        //构造系统 prompt
        String systemPrompt = """
            角色与目标：你是一个招聘助手，会针对用户的问题，结合候选人经历，岗位匹配度等专业知识，给用户提供指导。
            指导原则：你需要确保给出的建议合理科学，不会对候选人的表现有言论侮辱。
            限制：在提供建议时，需要强调在个性建议方面用户仍然需要线下寻求专业咨询。
            澄清：在与用户交互过程中，你需要明确回答用户关于招聘方面的问题，对于非招聘方面的问题，你的回应是‘我只是一个招聘助手，不能回答这个问题哦’。
            个性化：在回答时，你需要以专业可靠的预期回答，偶尔可以带点幽默感。调节气氛。
            给你提供一个数据参考，并且给你调用岗位投递检索公户
            请你跟进数据参考与工具返回结果回复用户的请求。
            """;
        //构造用户 prompt
        String userPrompt = """
            给你提供一些数据参考：{info},请用中文回答我的问题：{query}。
            请你跟进数据参考与工具返回结果回复用户的请求。
            """;
        //构造提示词
        //        String info = "刘磊是一个非常优秀的算法工程师，他可以参加面试";
        //        query = "刘磊可以参加面试吗";
        SystemMessage systemMessage = new SystemMessage(systemPrompt);
        PromptTemplate promptTemplate = new PromptTemplate(userPrompt);
        // 这里必须是info 和 query,参考的提示模板的关键词
        Message userMessage = promptTemplate.createMessage(Map.of("info", info, "query", query));
        OllamaOptions options = OllamaOptions.builder()
            //.withFunctions(Set.of("recruitServiceFunction"))
            .withModel("llama3.1")
            .build();

        Prompt prompt = new Prompt(List.of(userMessage  ), options);
        long currentTime3 = System.currentTimeMillis();
        //ChatResponse call = chatModel.call(prompt);
        Flux<String> stream = chatModel.stream(prompt.toString());
        log.info("模型调用耗时：{}ms", System.currentTimeMillis() - currentTime3);

        // 在消息的末尾添加来源信息
        Flux<String> messagesWithSource = stream.concatWith(Flux.just("来源："+fromFiles));
        return messagesWithSource;
        // 在所有消息的末尾添加来源信息
        //return "响应结果：" + call.getResult().getOutput().getContent() + ";来源：" + fromFiles;
    }

    @GetMapping("/wether")
    public String testFunction() {
        UserMessage userMessage = new UserMessage("济南天气怎么样？");
        OllamaOptions options = OllamaOptions.builder()
            .withFunction("locationWeatherFunction")
            .withModel("llama3.1")
            .build();
        Prompt prompt = new Prompt(List.of(userMessage), options);
        ChatResponse call = chatModel.call(prompt);
        return call.getResult().getOutput().getContent();
    }
}
